#include "llvm/Module.h"
#include "llvm/DerivedTypes.h"
#include "llvm/Instructions.h"
#include "llvm/Constants.h"
#include "llvm/BasicBlock.h"
#include "llvm/Pass.h"
#include "llvm/Support/CFG.h"
#include "llvm/Support/CallSite.h"
#include "llvm/IntrinsicInst.h"
#include "llvm/LinkAllVMCore.h"
#include "llvm/CodeGen/DwarfWriter.h"

#include "Passes.h"

using namespace std;
using namespace llvm;

// Heming: do not forget this!
//  kInstIDs = ICP0->getInstMap();

namespace klee {
	char InsertSubroutinePass::ID;

	bool InsertSubroutinePass::isCreateThreadFunction(Instruction *Inst) {
		Function *F;
		
		if (Inst->getOpcode() == Instruction::Call) {
			F = dyn_cast<CallInst>(Inst)->getCalledFunction();
			if (F && (F->getName() == "pthread_create" 
				//|| F->getName() == "pthread_create"
				|| false)) {
				return true;
			}
		}

		return false;
	}

bool InsertSubroutinePass::runOnModule(Module &M) {
	ObjectID &IDM = getAnalysis<ObjectID>();
	IDM.runOnModule(M);

	for (Module::iterator f = M.begin(), fe = M.end(); f != fe; ++f)
		for (Function::iterator b = f->begin(), be = f->end(); b != be; ++b)
			for (BasicBlock::iterator i = b->begin(), ie = b->end(); i != ie; ++i) {     
				Instruction *inst = i;
				instIDs[(long)inst] = IDM.getInstructionID(inst);
	}

	bool result = false;
	for (Module::iterator f = M.begin(), fe = M.end(); f != fe; ++f)
		for (Function::iterator b = f->begin(), be = f->end(); b != be; ++b)
			result |= runOnBasicBlock(*b);

	//fprintf(stderr, "finished InsertSubroutinePass::runOnModule\n");
	return result;
}


	bool InsertSubroutinePass::runOnBasicBlock(BasicBlock &BB) {
		bool result = false;
		
		for (BasicBlock::iterator IB = BB.begin(), IE = BB.end(); IB != IE; IB++) {
			if (isCreateThreadFunction(IB)) {
				// Insert two instructions, one is calling the subroutine directly, the other is exit().
				//Function *F = dyn_cast<CallInst>(IB)->getCalledFunction();
				//fprintf(stderr, "Function %s\n", F->getNameStr().c_str());
				std::vector<Value*> Args;
				CallSite CS = CallSite::get(IB);
				BasicBlock::iterator BI = IB;
				BI++;

				// insert pthread_self();
				std::vector<Value*> PthreadSelf_Args;
				StringRef PthreadSelf_SR("pthread_self");
				Function *PthreadSelf_DS = BB.getParent()->getParent()->getFunction(PthreadSelf_SR);
				assert(PthreadSelf_DS && "If there is no pthread_self() used in original program, please add \"pthread_self();\" to the end of the main() function.");
 				CallInst *CI = CallInst::Create(PthreadSelf_DS, PthreadSelf_Args.begin(), PthreadSelf_Args.end(), "", BI);

				// store pthread_self() to a place.
				Value *ThreadIdAddr = CS.getArgument(0);
				StoreInst *SI = new StoreInst(CI, ThreadIdAddr, BI);
				SI->setAlignment(4);

				// insert calling subrouting function.
				Value *SubRoutineName = CS.getArgument(2);
				Args.push_back(CS.getArgument(3));
				CallInst::Create(SubRoutineName, Args.begin(), Args.end(), "", BI);

				result = true;
			}
		}

		return result;
	}
}

